import cv2
import numpy as np
import pickle

cam = cv2.VideoCapture(2)

cv2.namedWindow("test")
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
img_counter = 0

objpoints = []
imgpoints = []

objp = np.zeros((6*7,3), np.float32)
objp[:,:2] = np.mgrid[0:7,0:6].T.reshape(-1,2)


while True:
    print(len(objpoints))
    rec, frame = cam.read()
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    ret, corners = cv2.findChessboardCorners(gray, (7, 6), None)
    if ret == True:
        objpoints.append(objp)

        corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
        imgpoints.append(corners2)

        frame = cv2.drawChessboardCorners(frame, (7, 6), corners2, ret)

    cv2.imshow("test", frame)
    if not rec:
        break

    k = cv2.waitKey(15)

    if k%256 == 27:
        # ESC pressed
        print("Escape hit, closing...")
        break

cam.release()

cv2.destroyAllWindows()

with open('objpoints','wb') as f:
    pickle.dump(objpoints, f)

with open('imgpoints','wb') as ff:
    pickle.dump(imgpoints, ff)


# ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1],None,None)
# h, w = frame.shape[:2]
# newcameramtx, roi=cv2.getOptimalNewCameraMatrix(mtx,dist,(w,h),1,(w,h))

# cv2.namedWindow("test1")
# cam = cv2.VideoCapture(2)
# while True:
#     rec, frame = cam.read()
#     dst = cv2.undistort(frame, mtx, dist, None, newcameramtx)
#     cv2.imshow("test", dst)
#
#     k = cv2.waitKey(15)
#
#     if k%256 == 27:
#         # ESC pressed
#         print("Escape hit, closing...")
#         break